import cv2
import numpy as np

class DCT:
    def __init__(self):
        self.block_size = 8
        self.terminator = "11111111111111100000000000000000"  # 32位结束标记，更安全
        self.embed_positions = [(3, 3), (3, 4), (4, 3), (4, 4)]  # 多个中频位置
        
    def embed(self, image_path, message, output_path):
        try:
            # 读取图像
            img = cv2.imread(image_path)
            if img is None:
                print(f"无法读取图像: {image_path}")
                return False
            
            print(f"图像形状: {img.shape}")
            print(f"原始消息: {message}")
            print(f"消息类型: {type(message)}")
            
            # 转换为YUV色彩空间，在Y分量中嵌入
            img_yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)
            y_channel = img_yuv[:,:,0].astype(np.float32)
            
            # 将消息转换为UTF-8字节
            message_bytes = message.encode('utf-8')
            message_length = len(message_bytes)
            
            print(f"UTF-8字节长度: {message_length}")
            print(f"UTF-8字节: {list(message_bytes)}")
            
            # 构建完整数据：长度(4字节) + 消息数据 + 结束标记
            length_bytes = message_length.to_bytes(4, byteorder='big')  # 4字节长度
            full_data = length_bytes + message_bytes
            
            # 转换为二进制
            full_binary = ''.join([format(byte, '08b') for byte in full_data])
            full_binary += self.terminator
            
            print(f"总二进制长度: {len(full_binary)} 位")
            print(f"长度信息: {length_bytes.hex()} ({message_length} 字节)")
            
            # 检查容量
            bits_per_block = len(self.embed_positions)
            blocks_needed = (len(full_binary) + bits_per_block - 1) // bits_per_block  # 向上取整
                
            available_blocks = (y_channel.shape[0] // self.block_size) * (y_channel.shape[1] // self.block_size)
            
            if blocks_needed > available_blocks:
                print(f"需要 {blocks_needed} 个块，但只有 {available_blocks} 个可用")
                return False
            
            print(f"可用块数: {available_blocks}, 需要块数: {blocks_needed}")
            
            # 嵌入数据
            bit_index = 0
            height, width = y_channel.shape
            
            blocks_used = 0
            for i in range(0, height - self.block_size, self.block_size):
                for j in range(0, width - self.block_size, self.block_size):
                    if bit_index >= len(full_binary):
                        break
                    
                    blocks_used += 1
                    # 提取8x8块
                    block = y_channel[i:i+self.block_size, j:j+self.block_size]
                    
                    # DCT变换
                    dct_block = cv2.dct(block)
                    
                    # 在多个位置嵌入数据
                    for pos_idx, (x, y) in enumerate(self.embed_positions):
                        if bit_index < len(full_binary):
                            bit = int(full_binary[bit_index])
                            
                            # 根据bit值调整系数
                            current_value = dct_block[x, y]
                            threshold = 25.0  # 提高阈值确保稳定性
                            
                            if bit == 1:
                                new_value = max(abs(current_value), threshold)
                                dct_block[x, y] = new_value
                            else:
                                new_value = -max(abs(current_value), threshold)
                                dct_block[x, y] = new_value
                                
                            bit_index += 1
                    
                    # 逆DCT变换
                    idct_block = cv2.idct(dct_block)
                    
                    # 放回原位置
                    y_channel[i:i+self.block_size, j:j+self.block_size] = idct_block
                
                if bit_index >= len(full_binary):
                    break
            
            print(f"成功嵌入 {bit_index} 位数据，使用 {blocks_used} 个块")
            
            # 转换回BGR
            img_yuv[:,:,0] = np.clip(y_channel, 0, 255)
            img_embedded = cv2.cvtColor(img_yuv.astype(np.uint8), cv2.COLOR_YUV2BGR)
            
            # 保存为PNG格式避免压缩
            if output_path.lower().endswith('.jpg') or output_path.lower().endswith('.jpeg'):
                output_path = output_path.rsplit('.', 1)[0] + '.png'
            
            cv2.imwrite(output_path, img_embedded, [cv2.IMWRITE_PNG_COMPRESSION, 0])
            print(f"图像保存成功: {output_path}")
            return True
            
        except Exception as e:
            print(f"DCT嵌入错误: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    def extract(self, image_path):
        try:
            # 读取图像
            img = cv2.imread(image_path)
            if img is None:
                print(f"无法读取图像: {image_path}")
                return ""
            
            print(f"提取图像形状: {img.shape}")
            
            # 转换为YUV色彩空间
            img_yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)
            y_channel = img_yuv[:,:,0].astype(np.float32)
            
            # 提取二进制数据
            binary_data = ""
            height, width = y_channel.shape
            
            # 提取足够的数据
            max_blocks_to_check = 10000  # 限制检查的块数
            
            blocks_checked = 0
            for i in range(0, height - self.block_size, self.block_size):
                for j in range(0, width - self.block_size, self.block_size):
                    if blocks_checked >= max_blocks_to_check:
                        break
                    
                    blocks_checked += 1
                    # 提取8x8块
                    block = y_channel[i:i+self.block_size, j:j+self.block_size]
                    
                    # DCT变换
                    dct_block = cv2.dct(block)
                    
                    # 从多个位置提取数据
                    for x, y in self.embed_positions:
                        coeff = dct_block[x, y]
                        # 根据系数符号提取bit，使用绝对值阈值
                        bit = '1' if coeff > 10.0 else '0'  # 提高阈值减少误判
                        binary_data += bit
            
            print(f"检查了 {blocks_checked} 个块")
            print(f"提取的二进制数据长度: {len(binary_data)} 位")
            
            # 查找结束标记
            terminator_pos = binary_data.find(self.terminator)
            if terminator_pos == -1:
                print("未找到结束标记")
                return "未找到结束标记"
            
            print(f"找到结束标记在位置: {terminator_pos}")
            
            # 提取有效数据（不包括结束标记）
            valid_binary = binary_data[:terminator_pos]
            print(f"有效数据长度: {len(valid_binary)} 位")
            
            # 检查最小长度
            if len(valid_binary) < 32:  # 至少需要4字节长度信息
                return "数据长度不足"
            
            # 提取长度信息（前4字节 = 32位）
            length_binary = valid_binary[:32]
            try:
                message_length = int(length_binary, 2)
                print(f"消息字节长度: {message_length}")
            except:
                return "长度信息无效"
            
            # 计算需要的总位数
            total_bits_needed = 32 + (message_length * 8)  # 长度(32位) + 消息数据
            
            if len(valid_binary) < total_bits_needed:
                return f"数据不完整，需要 {total_bits_needed} 位，只有 {len(valid_binary)} 位"
            
            # 提取消息数据
            message_binary = valid_binary[32:32 + (message_length * 8)]
            
            # 转换为字节
            message_bytes = bytearray()
            for i in range(0, len(message_binary), 8):
                byte_str = message_binary[i:i+8]
                if len(byte_str) == 8:
                    try:
                        byte_val = int(byte_str, 2)
                        message_bytes.append(byte_val)
                    except:
                        continue
            
            print(f"提取的字节数: {len(message_bytes)}")
            print(f"提取的字节: {list(message_bytes)}")
            
            # 尝试UTF-8解码
            try:
                result = message_bytes.decode('utf-8')
                print(f"成功解码消息: {result}")
                return result
            except UnicodeDecodeError as e:
                print(f"UTF-8解码失败: {e}")
                # 尝试其他编码
                try:
                    result = message_bytes.decode('gbk')
                    print(f"使用GBK解码: {result}")
                    return result
                except:
                    # 返回十六进制表示
                    hex_representation = message_bytes.hex()
                    print(f"原始字节(十六进制): {hex_representation}")
                    return f"解码失败，原始数据: {hex_representation}"
            
        except Exception as e:
            print(f"DCT提取错误: {e}")
            import traceback
            traceback.print_exc()
            return f"提取错误: {str(e)}"